"""
Created on 23-Jul-2009

@author: dexen
"""

from dexen.system.constants import NS_NAME

import os
import sys
import time

import rpyc
from rpyc.utils import server

from custom_service import MyService

from multiprocessing import Process, Queue, Manager, Lock

from dexen.system.server import queue_manager, pop_manager
from dexen.system import util

#===============================================================================
# GLOBALS
#===============================================================================
SLAVE_QUERY_INTERVAL = 60 # SECS

#===============================================================================
# SYNCHRONIZED SHARED DATA STRUCTURES 
#===============================================================================
queues = {} 
qlocks = {}
inds_db = Manager().dict() 
ind_db_lock = Lock()

#===============================================================================
# HELPER FUNCTIONS
#===============================================================================
def job_dispatcher(schema_name, schema_zip, pop_name, scripts_info, queues, 
                   slave_name):
    print "I am in job dispatcher"
    conn = rpyc.connect_by_service(service_name=slave_name, config={
                                        "allow_pickle":True,
                                        "instantiate_custom_exceptions":True})
    
    sys.path.append(os.path.join(os.getcwd(), "temp")) #TODO: check
    os.chdir(os.path.join(os.getcwd(), "temp"))
    slave = conn.root
    slave.setup_schema(schema_zip, schema_name, scripts_info)
    slave_attrs = slave.get_attributes()
    pop_mgr = pop_manager.PopManager(inds_db, ind_db_lock, scripts_info)
    queue_mgr = queue_manager.QManager(queues, qlocks, scripts_info, pop_mgr)
    
    while True:
        print queue_mgr
        script_name, input = queue_mgr.get_inds()
        print "executing ", script_name
        try:
            output = slave.execute_script(script_name, input)
            output = rpyc.classic.obtain(output)
            print "finished"            
        except Exception, e:
            print e
            print "Slave has an exception I am recovering and exiting!"
            import traceback
            traceback.print_tb(sys.exc_info()[2])
            queue_mgr.recover_inds(input, script_name)
            break
        queue_mgr.set_inds(script_name, output)

def slave_manager(schema_name, schema_zip, pop_name, scripts_info):
    """Administer connecting/disconnecting slaves.
    
    """
    conn = rpyc.connect_by_service(service_name=NS_NAME)
    naming_server = conn.root
    procs = []
    while True:
        slave_names = naming_server.get_free_slave_names()
        print "free slaves", slave_names
        for slave_name in slave_names:
            proc = Process(target=job_dispatcher, args=(schema_name, schema_zip,
                                                        pop_name, scripts_info,
                                                        queues, slave_name)
            )
            proc.daemonic = True
            procs.append(proc)
            proc.start()
        time.sleep(SLAVE_QUERY_INTERVAL)


#===============================================================================
# JOB SERVICE CLASS
#===============================================================================
class JobService(MyService):
    ALIASES = ["JOBSERVER"]
    
    def on_connect(self):
        MyService.on_connect(self)
    
    def on_disconnect(self):
        MyService.on_disconnect(self)
    
    def _get_script_obj(self, file_name, class_name):
        module = __import__(file_name, fromlist=[""])
        classobj = getattr(module, class_name)    
        return classobj()
    
    def _execute_init(self, schema_name, scripts_info):
        prev_dir = os.getcwd()
        os.chdir(os.path.join(prev_dir, "temp"))   
    
        for script_name, item in scripts_info.items():
            obj = self._get_script_obj(item["file_name"], item["class_name"])
            if obj.__isinit__():
                print script_name, "is initdir"
                obj.__initialize__()
                res = obj.__execute__()
                pop_mgr = pop_manager.PopManager(inds_db, ind_db_lock, 
                                                 scripts_info)
                q_mgr = queue_manager.QManager(queues, qlocks,scripts_info, 
                                               pop_mgr)
                q_mgr.set_inds(script_name, res)
        os.chdir(prev_dir)
        
    def _setup_schema(self, schema_name):
        schema_zip, modified_time = mserver.get_schema(schema_name)        
        temp_dir = os.path.abspath("temp")
        sys.path.append(temp_dir)
        util.unzip(schema_zip, temp_dir) # TODO: consider folder structure
        return schema_zip
        
    def _setup_queues(self, schema_name, pop_name):
        """Initialize queues and their associated locks.
        
        After initialization of data structures, init script has been executed
        to kick-off the run. scripts_info is retrieved from the main database.
        
        """
        global queues, qlocks

        scripts_info = mserver.get_scripts_info(schema_name, pop_name)
        scripts_info = rpyc.classic.obtain(scripts_info)
    
        for script_name, item in scripts_info.items():
            obj = self._get_script_obj(item["file_name"], item["class_name"])
            item["obj"] = obj
            if obj.__hasinput__():
                queues[script_name] = Queue()
                qlocks[script_name] = Lock()
        
        self._execute_init(schema_name, scripts_info)
                
        print scripts_info, "*"*23   
        return scripts_info
    
    def run_population(self, schema_name, pop_name):
        schema_zip = self._setup_schema(schema_name)
        scripts_info = self._setup_queues(schema_name, pop_name)
                 
        self.slave_mgr = Process(target=slave_manager, args=(schema_name,
                                                             schema_zip,
                                                             pop_name,
                                                             scripts_info
                                                            )
        )
        self.slave_mgr.daemonic = True
        self.slave_mgr.start()
          
    def stop_population(self):
        pass
    
    
def main():
    global db_conn, mserver
    db_conn = rpyc.connect_by_service(service_name="FRONTEND")
    mserver = db_conn.root
    jserver = server.ThreadedServer(JobService, 
                                protocol_config={"allow_pickle":True})
    jserver.start()

if __name__ == '__main__':
    main()
